from typing import Callable, Generic, Optional, TypeVar, Union, overload

from zkl_pyutils_registries.basics import MutableRegistry, Registry

AnyKey = TypeVar("AnyKey")
AnyQuery = TypeVar("AnyQuery")
AnyCallable = TypeVar("AnyCallable", bound=Callable)


class CallablesRegistry(Registry[AnyQuery, AnyCallable]):
    def __init__(self, base: Registry[AnyQuery, AnyCallable]):
        self._base = base

    def match(self, query: AnyQuery) -> Union[AnyCallable, 'MatchedCallables[AnyCallable]']:
        return MatchedCallables(self._base.match(query))


class MutableCallablesRegistry(
    CallablesRegistry[AnyQuery, AnyCallable],
    MutableRegistry[AnyQuery, AnyKey, AnyCallable],
    Generic[AnyQuery, AnyKey, AnyCallable],
):
    def __init__(self, base: MutableRegistry[AnyQuery, AnyKey, AnyCallable]):
        super().__init__(base)

    # noinspection PyMethodOverriding
    @overload
    def register(self, key: AnyKey) -> Callable[[AnyCallable], AnyCallable]:
        pass

    @overload
    def register(self, key: AnyKey, value: AnyCallable) -> AnyCallable:
        pass

    def register(self,
        key: Union[AnyKey, AnyCallable],
        value: Optional[AnyCallable] = None,
    ) -> Union[Callable[[AnyCallable], AnyCallable], AnyCallable]:
        if value is None:
            def decorated_func(value: AnyCallable) -> AnyCallable:
                self.register(key, value)
                return value

            return decorated_func

        assert isinstance(self._base, MutableRegistry)
        self._base.register(key, value)
        return value


class MatchedCallables(tuple[AnyCallable, ...], Callable):
    def __call__(self, *args, **kwargs):
        callables_error = []
        for callable_ in self:
            try:
                return callable_(*args, **kwargs)
            except NotImplementedError as error:
                callables_error.append(error)

        if not callables_error:
            raise NotImplementedError(f"No matched callable.")

        if len(callables_error) == 1:
            raise callables_error[0]

        raise NotImplementedErrorGroup(f"Failed to call all impls.", callables_error)


class NotImplementedErrorGroup(ExceptionGroup, NotImplementedError):
    pass
